//
//  MetalAdder.m
//  Metal-Master
//
//  Created by 杨冰冰 on 2020/6/15.
//  Copyright © 2020 杨冰冰. All rights reserved.
//

#import "MetalAdder.h"
#define ArrayLength 4
@implementation MetalAdder

- (instancetype)initWithDevice:(id <MTLDevice>)device {
    self = [super init];
    if (self) {
        //device
        _mDevice = device;
        
        NSError *error = nil;
        //library 从device中获取默认的着色器库
        id <MTLLibrary> defaultLibrary = [_mDevice newDefaultLibrary];
        if (defaultLibrary == nil) {
            NSLog(@"Failed to find the default library.");
            return nil;
        }
        //取着色器
        id <MTLFunction> addFunction = [defaultLibrary newFunctionWithName:@"add_arrays"];
        if (addFunction == nil) {
            NSLog(@"Failed to find the adder function.");
            return nil;
        }
        //获取管道状态对象
        _mAddFunctionPSO = [_mDevice newComputePipelineStateWithFunction:addFunction error:&error];
        
        //获取命令队列
        _commandQueue = [_mDevice newCommandQueue];
        
        //获取commandBuffer
        id <MTLCommandBuffer> commandBuffer = [_commandQueue commandBuffer];
        //创建缓冲
        NSUInteger bufferSize = ArrayLength * sizeof(float);
        _mBufferA = [_mDevice newBufferWithLength:bufferSize options:MTLResourceStorageModeShared];
        _mBufferB = [_mDevice newBufferWithLength:bufferSize options:MTLResourceStorageModeShared];
        _mBufferResult = [_mDevice newBufferWithLength:bufferSize options:MTLResourceStorageModeShared];
        
        [self generateRandomFloatData:_mBufferA];
        [self generateRandomFloatData:_mBufferB];
        
        //编码缓冲到command buffer
        id <MTLComputeCommandEncoder> computeEncoder = [commandBuffer computeCommandEncoder];
        [computeEncoder setComputePipelineState:_mAddFunctionPSO];
        [computeEncoder setBuffer:_mBufferA offset:0 atIndex:0];
        [computeEncoder setBuffer:_mBufferB offset:0 atIndex:1];
        [computeEncoder setBuffer:_mBufferResult offset:0 atIndex:2];
        
        //指定编码的线程数量和线程组大小
        MTLSize gridSize = MTLSizeMake(ArrayLength, 1, 1);
        
        NSUInteger threadGroupSize = _mAddFunctionPSO.maxTotalThreadsPerThreadgroup;
        if (threadGroupSize > ArrayLength) {
            threadGroupSize = ArrayLength;
        }
        MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1);

        //执行编码
        [computeEncoder dispatchThreadgroups:gridSize threadsPerThreadgroup:threadgroupSize];
        //结束buffer内编码
        [computeEncoder endEncoding];
        
        //提交一个buffer
        [commandBuffer commit];
        
        //等待执行结束。也可以[commandBuffer addCompletedHandler:]
        [commandBuffer waitUntilCompleted];
        
    }
    return self;
}

- (void)verifyResultes {
    float *a = _mBufferA.contents;
    float *b = _mBufferB.contents;
    float *result = _mBufferResult.contents;
    
    for (unsigned long index = 0; index <ArrayLength; index++) {
        assert(result[index] == a[index] + b[index]);
        printf("%f\n",result[index]);
    }
}

- (void)generateRandomFloatData:(id <MTLBuffer>)buffer {
    float *dataPtr = buffer.contents;
    
    for (unsigned long index = 0; index < ArrayLength; index++) {
        dataPtr[index] = (float)rand()/(float)(RAND_MAX);
    }
}
@end
